import numpy as np

# 根据分类数产生对应颜色
def bit_get(val, idx):
    """Gets the bit value.
    Args:
      val: Input value, int or numpy int array.
      idx: Which bit of the input val.
    Returns:
      The "idx"-th bit of input val.
    """
    return (val >> idx) & 1

def create_pascal_label_colormap(class_num):
    """Creates a label colormap used in PASCAL VOC segmentation benchmark.
    Returns:
      A colormap for visualizing segmentation results.
    """
    colormap = np.zeros((class_num, 3), dtype=int)
    ind = np.arange(class_num, dtype=int)

    for shift in reversed(range(8)):
        for channel in range(3):
            colormap[:, channel] |= bit_get(ind, channel) << shift
        ind >>= 3

    return colormap

# 用于显示语义分割结果
def to_color_img(img, color_list):
    score_i = img.argmax(dim=1).squeeze(dim=0)
    score_i = score_i.cpu().numpy()
    color_img = color_list[score_i].astype(np.uint8)
    return color_img
    
if __name__ == "__main__":
    color = create_pascal_label_colormap(34)
    print(color)